import argparse
import torchaudio
import shutil
import uuid
from fastapi import FastAPI, Form, Request
from fastapi.responses import JSONResponse
from typing import Optional
import os
import httpx
from pydub import AudioSegment
from config import USER_VOICE_FOLDER, TEMP_USER_VOICE_FOLDER, prompt_sr, GO_SERVICE_SECRET_KEY, ALGORITHM, REDIS_HOST, REDIS_DB
import redis
from tools.utils import is_audio_duration_within_limit, generate_timbre_id, decode_verify_token
from config import demo_audio_dir, configure_logging, clone_task_type, RANDOM_SEED, DEMO_AUDIO_OSS_DIR, timbre_clone_port
from loguru import logger
from tools._oss import oss_client
from task_queue2 import Task, TaskManager, start_scheduled_task
from server.db.repository.voice_clone_task_repository import add_voice_clone_task_to_db
from server.db.base import Base, engine
from fastapi import HTTPException
configure_logging()  # 配置日志

# 如果数据表不存在，则创建数据表
Base.metadata.create_all(bind=engine)

# # 将第三方库路径添加到系统路径中
# sys.path.append('{}/third_party/Matcha-TTS'.format(ROOT_DIR))

# 初始化FastAPI应用，提供详细的API描述
app = FastAPI(title="Mars AI", description="MarsX AI TTS&CLONE API", version="0.0.1")

# 在应用启动时初始化Redis客户端（建议放在模块顶层）
redis_pool = redis.ConnectionPool(
    host=REDIS_HOST,
    port=6379,
    db=REDIS_DB,
    max_connections=10,
    socket_keepalive=True,
    retry_on_timeout=True,
    socket_timeout=10
)
global_redis_client = redis.Redis(
    connection_pool=redis_pool,
    health_check_interval=30,
    socket_keepalive=True
)

# 使用时传入redis客户端实例
task_manager = TaskManager(max_concurrent_tasks=1, redis_client=global_redis_client)

# 启动定时任务调度
start_scheduled_task(task_manager)

# 初始化命令行参数解析器，支持服务器的自定义主机和端口
parser = argparse.ArgumentParser()
parser.add_argument('--host', type=str, default='0.0.0.0')
parser.add_argument('--port', type=int, default=8001)
args = parser.parse_args()


# API: 训练音色克隆任务
@app.post("/train", summary="音色克隆")
async def clone_voice(uid: str = Form(..., description="用户ID"),
                      timbre_id: Optional[str] = Form(None, description="音色ID"),
                      url: str = Form(..., description="音频文件地址"),
                      task_id: Optional[str] = Form(None, description="任务ID"),
                      seed: Optional[int] = Form(RANDOM_SEED, description="随机种子")
                      ):
    """
    克隆音色，并后台处理音频生成任务。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param url: 音频文件地址
    :param task_id: 任务ID
    :param seed: 随机种子
    :param background_tasks: FastAPI后台任务处理对象
    :return: JSON 响应
    """
    voice_id = timbre_id
    # 生成音色ID
    if not timbre_id:
        voice_id = generate_timbre_id()
        timbre_id = f"{voice_id}"
        logger.debug(f"自己生成的INT32音色ID：{voice_id}, 转换成字符串的音色ID：{timbre_id}")
    else:
        logger.debug(f"使用传入的音色ID：{timbre_id}")
    # 生成任务ID
    if not task_id:
        task_id = f"t_{str(uuid.uuid4())}"
        logger.debug(f"自己生成的任务ID：{task_id}")
    else:
        logger.debug(f"使用传入的任务ID：{task_id}")

    # 检查URL的有效性
    if not oss_client.check_url_validity(url):
        logger.error(f"无效的文件地址或URL不可访问：{url}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "无效的文件地址或URL不可访问，请提供有效的 HTTP/HTTPS 地址。",
                                                      "task_id": task_id})

    # 定义临时目录和本地文件路径
    temp_directory = os.path.join(TEMP_USER_VOICE_FOLDER, uid, timbre_id)
    object_name = os.path.basename(url)
    local_file_path = os.path.join(temp_directory, object_name)

    os.makedirs(temp_directory, exist_ok=True)  # 创建临时目录

    # 下载音频文件并保存到本地
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
            response.raise_for_status()  # 检查响应状态码是否为200
            with open(local_file_path, 'wb') as f:
                f.write(response.content)
        logger.debug(f"成功下载文件：{object_name}")
    except Exception as e:
        logger.error(f"下载音频文件失败：{str(e)}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "下载音频文件失败，请检查文件地址是否正确。",
                                                      "task_id": task_id})
    try:
        audio = AudioSegment.from_file(local_file_path)

        audio.export(local_file_path, format="wav")
    except Exception as e:
        logger.error(f"音频文件转换失败：{str(e)}")
        shutil.rmtree(temp_directory)  # 删除临时目录
        logger.debug(f"删除临时目录：{temp_directory}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "音频文件转换失败，请检查文件格式是否正确。",
                                                      "task_id": task_id})
    # 验证音频文件时长是否在限制范围内（小于30秒）
    if not is_audio_duration_within_limit(local_file_path):
        logger.error(f"音频文件时长超过30秒。")
        shutil.rmtree(temp_directory)  # 删除临时目录
        logger.debug(f"删除临时目录：{temp_directory}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "音频文件时长超过30秒，请提供时长不超过30秒的音频文件。",
                                                      "task_id": task_id})

    # 验证音频文件的采样率
    if torchaudio.info(local_file_path).sample_rate < prompt_sr:
        logger.error(f"音频文件采样率过低：{torchaudio.info(local_file_path).sample_rate}")
        shutil.rmtree(temp_directory)  # 删除临时目录
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": f"音频文件采样率过低，请提供采样率不低于 {prompt_sr} kHz 的音频文件。",
                                                      "task_id": task_id})

    # 移动文件到用户目录
    user_directory = os.path.join(USER_VOICE_FOLDER, uid, timbre_id)

    final_file_path = os.path.join(user_directory, object_name)

    try:
        if os.path.exists(user_directory):
            shutil.rmtree(user_directory)
        os.makedirs(user_directory, exist_ok=True)  # 创建用户目录
        shutil.move(local_file_path, final_file_path)
        shutil.rmtree(temp_directory)  # 删除临时目录
        logger.debug(f"文件移动成功：{final_file_path}")
    except Exception as e:
        logger.error(f"移动文件失败：{str(e)}")
        task_manager.fail_task(task_id)
        logger.error(f"任务 {task_id} 执行失败,已标记为失败任务,并从数据库和Redis中删除任务")
        return JSONResponse(status_code=500,
                            content={"status_code": 500, "isSuccess": False, "message": "文件移动失败，请稍后重试。",
                                     "task_id": task_id})
    # 插入任务信息到数据库，初始状态为 'pending'
    try:
        add_voice_clone_task_to_db(uid, timbre_id, url, task_id, task_status='pending')
    except Exception as e:
        logger.error(f"插入任务信息到数据库失败：{str(e)}")
        shutil.rmtree(user_directory)  # 删除用户音色目录
        logger.debug(f"删除用户音色目录：{user_directory}")
        return JSONResponse(status_code=500, content={"status_code": 500, "isSuccess": False,
                                                      "message": "插入任务信息到数据库失败。",
                                                      "task_id": task_id})
    # 后台处理试听音频的生成任务
    try:
        task = Task(task_id=task_id, task_type=clone_task_type, uid=uid, timbre_id=timbre_id, seed=seed)
        task_manager.add_task(task)

        # 获取任务排队信息
        task_status = task_manager.check_task_status(task_id)

    except Exception as e:
        logger.error(f"生成试听音频时出错: {str(e)}")
        shutil.rmtree(user_directory)  # 删除用户音色目录
        logger.debug(f"删除用户音色目录：{user_directory}")
        try:
            task_manager.fail_task(task_id)
        except Exception as e:
            logger.error(f"更新任务状态失败：{str(e)}")
        return JSONResponse(status_code=500,
                            content={"status_code": 500, "message": "生成试听音频时出错。", "task_id": task_id})

    return JSONResponse(status_code=200,
                        content={"status_code": 200, "isSuccess": True, "message": "任务提交成功...",
                                 "timbre_id": voice_id, "task_id": task_id, "position": task_status['position'],
                                 "expected_wait_time": task_status['wait_time']})

# API: 训练音色克隆任务
@app.post("/v2/voice_clone", summary="音色克隆")
async def voice_clone(uid: str = Form(..., description="用户ID"),
                      timbre_id: Optional[str] = Form(None, description="音色ID"),
                      url: str = Form(..., description="音频文件地址"),
                      task_id: Optional[str] = Form(None, description="任务ID"),
                      seed: Optional[int] = Form(RANDOM_SEED, description="随机种子"),
                      request: Request = None
                      ):
    """
    克隆音色，并后台处理音频生成任务。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param url: 音频文件地址
    :param task_id: 任务ID
    :param seed: 随机种子
    :param background_tasks: FastAPI后台任务处理对象
    :return: JSON 响应
    """
    # 验证token
    try:
        token_info = decode_verify_token(request, GO_SERVICE_SECRET_KEY, ALGORITHM)
    except HTTPException as e:
        return JSONResponse(status_code=401, content={"status_code": 401, "isSuccess": False, "message": e.detail.get("msg"), "task_id": task_id})
    voice_id = timbre_id
    # 生成音色ID
    if not timbre_id:
        voice_id = generate_timbre_id()
        timbre_id = f"{voice_id}"
        logger.debug(f"自己生成的INT32音色ID：{voice_id}, 转换成字符串的音色ID：{timbre_id}")
    else:
        logger.debug(f"使用传入的音色ID：{timbre_id}")
    # 生成任务ID
    if not task_id:
        task_id = f"t_{str(uuid.uuid4())}"
        logger.debug(f"自己生成的任务ID：{task_id}")
    else:
        logger.debug(f"使用传入的任务ID：{task_id}")

    # 检查URL的有效性
    if not oss_client.check_url_validity(url):
        logger.error(f"无效的文件地址或URL不可访问：{url}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "无效的文件地址或URL不可访问，请提供有效的 HTTP/HTTPS 地址。",
                                                      "task_id": task_id})

    # 定义临时目录和本地文件路径
    temp_directory = os.path.join(TEMP_USER_VOICE_FOLDER, uid, timbre_id)
    object_name = os.path.basename(url)
    local_file_path = os.path.join(temp_directory, object_name)

    os.makedirs(temp_directory, exist_ok=True)  # 创建临时目录

    # 下载音频文件并保存到本地
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
            response.raise_for_status()  # 检查响应状态码是否为200
            with open(local_file_path, 'wb') as f:
                f.write(response.content)
        logger.debug(f"成功下载文件：{object_name}")
    except Exception as e:
        logger.error(f"下载音频文件失败：{str(e)}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "下载音频文件失败，请检查文件地址是否正确。",
                                                      "task_id": task_id})
    try:
        audio = AudioSegment.from_file(local_file_path)

        audio.export(local_file_path, format="wav")
    except Exception as e:
        logger.error(f"音频文件转换失败：{str(e)}")
        shutil.rmtree(temp_directory)  # 删除临时目录
        logger.debug(f"删除临时目录：{temp_directory}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "音频文件转换失败，请检查文件格式是否正确。",
                                                      "task_id": task_id})
    # 验证音频文件时长是否在限制范围内（小于30秒）
    if not is_audio_duration_within_limit(local_file_path):
        logger.error(f"音频文件时长超过30秒。")
        shutil.rmtree(temp_directory)  # 删除临时目录
        logger.debug(f"删除临时目录：{temp_directory}")
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": "音频文件时长超过30秒，请提供时长不超过30秒的音频文件。",
                                                      "task_id": task_id})

    # 验证音频文件的采样率
    if torchaudio.info(local_file_path).sample_rate < prompt_sr:
        logger.error(f"音频文件采样率过低：{torchaudio.info(local_file_path).sample_rate}")
        shutil.rmtree(temp_directory)  # 删除临时目录
        return JSONResponse(status_code=400, content={"status_code": 400, "isSuccess": False,
                                                      "message": f"音频文件采样率过低，请提供采样率不低于 {prompt_sr} kHz 的音频文件。",
                                                      "task_id": task_id})

    # 移动文件到用户目录
    user_directory = os.path.join(USER_VOICE_FOLDER, uid, timbre_id)

    final_file_path = os.path.join(user_directory, object_name)

    try:
        if os.path.exists(user_directory):
            shutil.rmtree(user_directory)
        os.makedirs(user_directory, exist_ok=True)  # 创建用户目录
        shutil.move(local_file_path, final_file_path)
        shutil.rmtree(temp_directory)  # 删除临时目录
        logger.debug(f"文件移动成功：{final_file_path}")
    except Exception as e:
        logger.error(f"移动文件失败：{str(e)}")
        task_manager.fail_task(task_id)
        logger.error(f"任务 {task_id} 执行失败,已标记为失败任务,并从数据库和Redis中删除任务")
        return JSONResponse(status_code=500,
                            content={"status_code": 500, "isSuccess": False, "message": "文件移动失败，请稍后重试。",
                                     "task_id": task_id})
    # 插入任务信息到数据库，初始状态为 'pending'
    try:
        add_voice_clone_task_to_db(uid, timbre_id, url, task_id, task_status='pending')
    except Exception as e:
        logger.error(f"插入任务信息到数据库失败：{str(e)}")
        shutil.rmtree(user_directory)  # 删除用户音色目录
        logger.debug(f"删除用户音色目录：{user_directory}")
        return JSONResponse(status_code=500, content={"status_code": 500, "isSuccess": False,
                                                      "message": "插入任务信息到数据库失败。",
                                                      "task_id": task_id})
    # 后台处理试听音频的生成任务
    try:
        task = Task(task_id=task_id, task_type=clone_task_type, uid=uid, timbre_id=timbre_id, seed=seed)
        task_manager.add_task(task)

        # 获取任务排队信息
        task_status = task_manager.check_task_status(task_id)

    except Exception as e:
        logger.error(f"生成试听音频时出错: {str(e)}")
        shutil.rmtree(user_directory)  # 删除用户音色目录
        logger.debug(f"删除用户音色目录：{user_directory}")
        try:
            task_manager.fail_task(task_id)
        except Exception as e:
            logger.error(f"更新任务状态失败：{str(e)}")
        return JSONResponse(status_code=500,
                            content={"status_code": 500, "message": "生成试听音频时出错。", "task_id": task_id})

    return JSONResponse(status_code=200,
                        content={"status_code": 200, "isSuccess": True, "message": "任务提交成功...",
                                 "timbre_id": voice_id, "task_id": task_id, "position": task_status['position'],
                                 "expected_wait_time": task_status['wait_time']})


# API: 检查音色克隆任务状态
@app.post("/check_file", summary="检查音色克隆任务是否完成")
async def check_file(uid: str = Form(..., description="用户ID"),
                     timbre_id: str = Form(..., description="音色ID"),
                     task_id: str = Form(..., description="任务ID")):
    """
    检查音色克隆任务是否完成，完成后返回试听文件URL，同时返回排队位置和预计等待时长。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param task_id: 任务ID
    :return: JSON 响应
    """
    task_status = task_manager.check_task_status(task_id)
    logger.debug(f"任务状态：{task_status}")

    if task_status['status'] == 'failed':
        logger.debug(f"UID: {uid}, 音色ID: {timbre_id}, 任务ID: {task_id} 的音色克隆任务失败...")
        return JSONResponse(status_code=200,
                            content={"status_code": 405, "isSuccess": False, "message": "音色克隆任务失败。",
                                     "task_id": task_id, "status": -1})
    if task_status['status'] == 'completed':
        # 获取试听音频文件的名称和OSS中的URL
        file_name = f"{uid}_{timbre_id}_demo.wav"
        preview_url = oss_client.get_file_url(file_name, directory=DEMO_AUDIO_OSS_DIR)

        # 检查OSS URL是否有效
        if oss_client.check_url_validity(preview_url):
            logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务已完成...")
            # 清理本地缓存
            local_audio_dir = os.path.join(demo_audio_dir, uid, timbre_id)
            if os.path.exists(local_audio_dir):
                shutil.rmtree(local_audio_dir)  # 清理本地生成的音频文件
            return JSONResponse(status_code=200,
                                content={"status_code": 200, "isSuccess": True, "message": "音色克隆任务已完成。",
                                         "preview_url": preview_url,
                                         "task_id": task_id,
                                         "timbre_id": timbre_id,
                                         "position": task_status.get('position', 0),  # 返回排队位置
                                         "expected_wait_time": task_status.get('wait_time', 0),
                                         "status": 1})  # 返回预计等待时长
        else:
            logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务尚未完成...")
            return JSONResponse(status_code=200,
                                content={"status_code": 200, "isSuccess": False, "message": "音色克隆任务尚未完成。",
                                         "task_id": task_id,
                                         "timbre_id": timbre_id,
                                         "position": task_status.get('position', 0),  # 返回排队位置
                                         "expected_wait_time": task_status.get('wait_time', 0),
                                         "status": -2})  # 返回预计等待时长
    if task_status['status'] == "not_found":
        logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务在数据库中未找到...")
        return JSONResponse(status_code=404,
                            content={"status_code": 404, "isSuccess": False,
                                     "message": f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务在数据库中未找到...",
                                     "task_id": task_id})
    else:
        logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务尚未完成...")
        return JSONResponse(status_code=200,
                            content={"status_code": 200, "isSuccess": False,
                                     "message": f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务尚未完成...",
                                     "task_id": task_id,
                                     "timbre_id": timbre_id,
                                     "position": task_status.get('position', 0),  # 返回排队位置
                                     "expected_wait_time": task_status.get('wait_time', 0),
                                     "status": -1})

# API: 检查音色克隆任务状态
@app.post("/v2/check_file", summary="检查音色克隆任务是否完成")
async def check_file_v2(uid: str = Form(..., description="用户ID"),
                     timbre_id: str = Form(..., description="音色ID"),
                     task_id: str = Form(..., description="任务ID"),
                     request: Request = None):
    """
    检查音色克隆任务是否完成，完成后返回试听文件URL，同时返回排队位置和预计等待时长。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param task_id: 任务ID
    :return: JSON 响应
    """
    try:
        token_info = decode_verify_token(request, GO_SERVICE_SECRET_KEY, ALGORITHM)
    except HTTPException as e:
        return JSONResponse(status_code=401, content={"status_code": 401, "isSuccess": False, "message": e.detail.get("msg"), "task_id": task_id})
    task_status = task_manager.check_task_status(task_id)
    logger.debug(f"任务状态：{task_status}")

    if task_status['status'] == 'failed':
        logger.debug(f"UID: {uid}, 音色ID: {timbre_id}, 任务ID: {task_id} 的音色克隆任务失败...")
        return JSONResponse(status_code=200,
                            content={"status_code": 405, "isSuccess": False, "message": "音色克隆任务失败。",
                                     "task_id": task_id, "status": -1})
    if task_status['status'] == 'completed':
        # 获取试听音频文件的名称和OSS中的URL
        file_name = f"{uid}_{timbre_id}_demo.wav"
        preview_url = oss_client.get_file_url(file_name, directory=DEMO_AUDIO_OSS_DIR)

        # 检查OSS URL是否有效
        if oss_client.check_url_validity(preview_url):
            logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务已完成...")
            # 清理本地缓存
            local_audio_dir = os.path.join(demo_audio_dir, uid, timbre_id)
            if os.path.exists(local_audio_dir):
                shutil.rmtree(local_audio_dir)  # 清理本地生成的音频文件
            return JSONResponse(status_code=200,
                                content={"status_code": 200, "isSuccess": True, "message": "音色克隆任务已完成。",
                                         "preview_url": preview_url,
                                         "task_id": task_id,
                                         "timbre_id": timbre_id,
                                         "position": task_status.get('position', 0),  # 返回排队位置
                                         "expected_wait_time": task_status.get('wait_time', 0),
                                         "status": 1})  # 返回预计等待时长
        else:
            logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务尚未完成...")
            return JSONResponse(status_code=200,
                                content={"status_code": 200, "isSuccess": False, "message": "音色克隆任务尚未完成。",
                                         "task_id": task_id,
                                         "timbre_id": timbre_id,
                                         "position": task_status.get('position', 0),  # 返回排队位置
                                         "expected_wait_time": task_status.get('wait_time', 0),
                                         "status": -2})  # 返回预计等待时长
    if task_status['status'] == "not_found":
        logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务在数据库中未找到...")
        return JSONResponse(status_code=404,
                            content={"status_code": 404, "isSuccess": False,
                                     "message": f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务在数据库中未找到...",
                                     "task_id": task_id})
    else:
        logger.debug(f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务尚未完成...")
        return JSONResponse(status_code=200,
                            content={"status_code": 200, "isSuccess": False,
                                     "message": f"UID: {uid}, 音色ID: {timbre_id} 的音色克隆任务尚未完成...",
                                     "task_id": task_id,
                                     "timbre_id": timbre_id,
                                     "position": task_status.get('position', 0),  # 返回排队位置
                                     "expected_wait_time": task_status.get('wait_time', 0),
                                     "status": -1})


# 启动应用
if __name__ == '__main__':
    import uvicorn
    # 运行FastAPI应用
    uvicorn.run(app, host="0.0.0.0", port=timbre_clone_port)
